from __future__ import division

from StringIO import StringIO

import numpy 

from paste.request import parse_dict_querystring
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib import rcParams
rcParams['xtick.labelsize'] = 'small'
rcParams['ytick.labelsize'] = 'small'

from pydap.model import *
from pydap.responses.lib import BaseResponse
from pydap.lib import walk


class PlotResponse(BaseResponse):

    __description__ = "Time series plot"

    def __init__(self, dataset):
        BaseResponse.__init__(self, dataset)
        self.headers.extend([
                ('Content-description', 'dods_plot'),
                ('Content-type', 'image/png'),
                ])

    def __call__(self, environ, start_response):
        self.serialize = self._get_plot(environ)
        return BaseResponse.__call__(self, environ, start_response)

    def _get_plot(self, environ):
        query = parse_dict_querystring(environ)
        dpi = float(environ.get('pydap.responses.plot.dpi', 80))

        def serialize(dataset):
            # Get 1D grids.
            grids = [g for g in walk(dataset, GridType)
                    if len(g.shape) == 1]
            n = len(grids)

            w = float(query.get('WIDTH', 400))
            h = float(query.get('HEIGHT', 400*n))
            figsize = w/dpi, h/dpi
            fig = Figure(figsize=figsize, dpi=dpi)
            fig.figurePatch.set_alpha(0.0)

            # Plot 1D grids.
            for i, grid in enumerate(grids):
                ax = fig.add_subplot(n, 1, i+1)
                y, x = grid.data
                y = numpy.asarray(y)
                x = numpy.asarray(x)
                if 'missing_value' in grid.attributes:
                    y = numpy.ma.masked_equal(y, grid.missing_value)
                elif '_FillValue' in grid.attributes:
                    y = numpy.ma.masked_equal(y, grid._FillValue)
                ax.plot(x, y, '-')
                ax.set_title(grid.name)
                ax.set_xlabel(('%s (%s)' % (
                    grid.dimensions[0], 
                    grid.maps.values()[0].attributes.get('units', ''))
                    ).replace(' ()', ''))
                ax.set_ylabel(grid.attributes.get('units', ''))

            # Save to buffer.
            canvas = FigureCanvas(fig)
            output = StringIO() 
            canvas.print_png(output)
            if hasattr(dataset, 'close'): dataset.close()
            return [ output.getvalue() ]
        return serialize
